import numpy as np


# 求解方程
def gaussian_elimination(A, b):
    """
    利用高斯消元法求方程组 Ax=b 的解
    :param A: 方程组系数矩阵,形状为 (n, n)
    :param b: 方程组常量矩阵，形状为 (n, )
    :return:方程组的解 x
    """
    # 将系数矩阵和右侧常数合并
    Ab = np.column_stack((A.astype(float), b.astype(float)))

    n = len(b)
    print(n)

    # # 将第一列元素进行消元，保留第一行第一列位置元素
    # # i 是要处理的行
    # for i in range(1, n):
    #     factor = Ab[i, 0] / Ab[0, 0]
    #     Ab[i, 0:] -= factor * Ab[0, 0:]
    #
    # # 将第二列元素进行消元，保留第二行第二列位置元素
    # for i in range(2, n):
    #     factor = Ab[i, 1] / Ab[0, 1]
    #     Ab[i, 1:] -= factor * Ab[0, 1:]

    # 针对存在 n 次不确定的情况，将上面重复的代码使用循环进行替换
    for j in range(0, n):
        for i in range(1 + j, n):
            factor = Ab[i, j] / Ab[0, j]
            Ab[i, j:] -= factor * Ab[0, j:]

    # 回代，计算 x 的值
    x = np.zeros(n)

    # x[2] = (Ab[2,3])/Ab[2,2]
    # x[1] = (Ab[1,3]-Ab[1,2]*x[2])/Ab[1,2]
    # x[0] =
    for i in range(n - 1, -1, -1):
        x[i] = (Ab[i, n] - np.dot(Ab[i, i:n], x[i:n])) / Ab[i, i]

    return x


if __name__ == '__main__':
    A = np.array(
        [
            [9, -1, -1],
            [-1, 8, 0],
            [-1, 0, 9]
        ]
    )

    b = np.array(
        [7, 7, 8]
    )

    x = gaussian_elimination(A, b)
    print(x)
